import torch
import torch.nn as nn
import torch.nn.functional as F

from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec
import timm


@BACKBONE_REGISTRY.register()
class SWSLResNet(Backbone):
    """
    DINOv2 model

    Args:
        model_name (str): The name of the model architecture 
            should be one of ('dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14')
        num_trainable_blocks (int): The number of last blocks in the model that are trainable.
        norm_layer (bool): If True, a normalization layer is applied in the forward pass.
        return_token (bool): If True, the forward pass returns both the feature map and the token.
    """
    def __init__( 
            self,
            cfg, input_shape
        ):
        super().__init__()
        self.backbone = timm.create_model('swsl_resnet18', features_only=True, output_stride=32,
                                          out_indices=(1, 2, 3, 4), pretrained=True)
        self._out_features = ["res2", "res3", "res4", "res5"]
        # encoder_channels = self.backbone.feature_info.channels()
        self._out_feature_strides = {
            "res2": 4,
            "res3": 8,
            "res4": 16,
            "res5": 32,
        }
        self._out_feature_channels = {
            "res2": 64,
            "res3": 128,
            "res4": 256,
            "res5": 512,
        }
    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
            )
            for name in self._out_features
        }
        
    def forward(self, x):
        """
        The forward method for the DINOv2 class

        Parameters:
            x (torch.Tensor): The input tensor [B, 3, H, W]. H and W should be divisible by 14.

        Returns:
            f (torch.Tensor): The feature map [B, C, H // 14, W // 14].
            t (torch.Tensor): The token [B, C]. This is only returned if return_token is True.
        """

        res2, res3, res4, res5 = self.backbone(x)
        
        outputs={
            "res2": res2,
            "res3": res3,
            "res4": res4,
            "res5": res5,
        }

        return outputs
